import numpy as np
import cv2
import os


def img2vector(imgfilename):
    img = cv2.imread(imgfilename, cv2.IMREAD_GRAYSCALE)
    rows, columns = img.shape
    #print(img.shape)
    vec = img.reshape(rows * columns)
    return vec


def distance(vector1, vector2):
    diff = vector1 - vector2  # 差
    double_diff = diff ** 2  # 差平方
    sum_of_double_diff = double_diff.sum()  # 差平方和
    dist = np.sqrt(sum_of_double_diff)
    return dist


# training_image_matrix是训练集图像数据矩阵，class_vector是类别向量
def classify(img, training_image_matrix, class_vector, k):
    N = len(training_image_matrix)  # 训练集的图像个数
    dist = np.zeros((N))  # 距离向量初始化为全 0
    # 计算待识别图像与训练集各个图像之间的欧几里德距离
    for i in range(N):
        dist[i] = distance(img, training_image_matrix[i])
    # 对距离进行排序
    sorted_index = dist.argsort() # numpy的间接排序，返回的是元素的下标
    match_count = {} # 用字典方式，便于数据成对
    # K 近邻   拿到前 K 个数值
    for i in range(k):
        match_class = class_vector[sorted_index[i]]
        match_count[match_class] = match_count.get(match_class,0) + 1 # 找不到则缺省为0

    # 对字典进行排序 从高到底 降序
    match_count_in_order = sorted(match_count.items(),key=lambda item:item[1],reverse=True)
    decided = match_count_in_order[0][0]
    return decided




training_dir = "../training"
sub_dir_and_files = os.listdir(training_dir)

sub_dirs = []
# 如果是目录
for x in sub_dir_and_files:
    if os.path.isdir(training_dir + "/" + x):
        sub_dirs.append(x)

# 计算装备图像总数
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(training_dir + "/" + subdir))

# 初始化训练图像数据矩阵 (N 行，128*128列) 和装备向量（长度为N）
training_img_matrix = np.zeros((N, 128 * 128))  # 每个图像一行数据
training_equipment_vector = [''] * N

i = 0  # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(training_dir + "/" + subdir)
    for image in image_files:
        # 将图像转换为向量
        v = img2vector(training_dir + "/" + subdir + "/" + image)
        training_img_matrix[i] = v
        training_equipment_vector[i] = subdir
        i += 1

# 通常采用测试用例的方式进行测试
testcases = ['bmp-2-1.jpg', 'btr-70-1.jpg', 't-72-1.jpg',
             't-72-3.jpg', 't-72-25.jpg']
img = 'bmp-2-1.jpg'

imgvector = img2vector(img)
recognized = classify(imgvector, training_img_matrix,
                      training_equipment_vector, 8)
print("识别结果为：",recognized)


